"""
@Time: 2020/12/9 下午 7:51
@Author: jinzhuan
@File: fn.py
@Desc: 
"""
import os
from cognlp import *


class FrameNetLoader:
    def __init__(self):
        self.label_set = set()

    def _load(self, path):
        datable = DataTable()
        sentence = []
        with open(path) as f:
            for line in f:
                if len(line) == 0 or line[0] == "\n":
                    if len(sentence) > 0:
                        datable('sentences', sentence)
                        datable('labels', label)
                        datable('positions', pos)
                        sentence = []
                    continue

                words = line.split('\t')
                sentence.append(words[1])
                if words[-3] not in '_':
                    label = words[-3]
                    pos = words[0]
                    self.label_set.add(words[-3])
            if len(sentence) > 0:
                datable('sentences', sentence)
                datable('labels', label)
                datable('positions', pos)
            if len(datable) == 0:
                raise RuntimeError("No data found {}.".format(path))
        return datable

    def load(self, path):
        train_set = self._load(os.path.join(path, 'train.bios'))
        dev_set = self._load(os.path.join(path, 'dev.bios'))
        test_set = self._load(os.path.join(path, 'test.bios'))
        return train_set, dev_set, test_set

    def get_labels(self):
        labels = list(self.label_set)
        labels.sort()
        return labels
